/**
 * Copyright (c) 2024 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

/*!
 * \file hccl.h
 * \brief
 */
#ifndef LIB_HCCL_HCCL_H
#define LIB_HCCL_HCCL_H

#include "hccl_common.h"
#include "../../impl/hccl/hccl_msg.h"

#if __CCE_AICORE__ == 220
namespace AscendC {
/*!
 * @class Hccl
 * @brief 本类主要对外提供了集合通信原语接口(下统称Prepare)，对标华为集合通信C++接口，包含AllReduce、AllGather、ReduceScatter
 *        本类典型的使用方式如下：
 *        1）用户创建Hccl对象，并初始化
 *        2）用户通过对应Prepare接口异步下发通信任务，服务端监听到后即开始通信任务的展开
 *        3）用户调用Commit接口通知执行对应任务，服务端方能执行对应通信任务
 *        4）用户调用Wait阻塞接口，等待服务端执行完对应的通信任务
 *        5）用户调用Finalize接口，通知服务端后续无通信任务，执行完即可退出
 * @tparam serverType: 当前只支持HcclServerType::HCCL_SERVER_TYPE_AICPU
 */
template <HcclServerType serverType = HcclServerType::HCCL_SERVER_TYPE_AICPU>
class Hccl {
public:
    /*!
     * @class Hccl
     * @brief 集合通信算子AllReduce的任务下发接口（任务内容为：将group内所有节点的同名张量进行reduce操作后，
     *        再把结果发送到所有节点的输出buffer），返回该任务的标识handleId给用户
     * @tparam commit: true/false 该参数设为true时，表示同步通知服务端可以执行对应的AllReduce通信任务
     * @param [out] recvBuf: 目的数据buffer地址，集合通信结果输出至此buffer中
     * @param [in] sendBuf: 源数据buffer地址
     * @param [in] count: 参与AllReduce操作的数据个数，比如只有一个int32数据参与，则count=1
     * @param [in] dataType: AllReduce操作的数据类型
     * @param [in] op: Reduce的操作类型，目前支持sum、prod、max、min等操作类型
     * @param [in] repeat: 表明本通信任务需要执行的次数，默认为1，repeat必须>=1。
     *                     当repeat>1时，偏移数据在服务端由count*sizeof(dataType)算出，且只能用于连续场景
     * @return 返回该任务的标识handleId给用户，handleId >= 0。当本接口调用失败时，返回 -1。
     */
    template <bool commit = false>
    __aicore__ inline HcclHandle AllReduce(
        GM_ADDR sendBuf, GM_ADDR recvBuf, uint64_t count, HcclDataType dataType, HcclReduceOp op, uint8_t repeat = 1);

    /*!
     * @class Hccl
     * @brief 集合通信算子AllGather的任务下发接口（任务内容为：将group内所有节点的输入按照rankid重新排序，然后拼接起来，
     *        再将结果发送到所有节点的输出)，返回该任务的标识handleId给用户
     * @tparam commit: true/false 该参数设为true时，表示同步通知服务端可以执行对应的AllGather通信任务
     * @param [out] recvBuf: 目的数据buffer地址，集合通信结果输出至此buffer中
     * @param [in] sendBuf: 源数据buffer地址
     * @param [in] sendCount: 参与AllGather操作的sendBuf的数据个数，recvBuf的数据size则等于sendCount * rank size。
     * @param [in] dataType: AllGather操作的数据类型
     * @param [in] strideCount: 当多个rank的数据拼接到一个rank时，由于可能是多轮切分，因此需要strideCount参数告诉每一轮切分的数据个数
     * @param [in] repeat: 表明本通信任务需要执行的次数，默认为1，repeat必须>=1。
     *                     当repeat>1时，偏移数据在服务端由count*sizeof(dataType)算出，且只能用于连续场景
     * @return 返回该任务的标识handleId给用户，handleId >= 0。当本接口调用失败时，返回 -1。
     */
    template <bool commit = false>
    __aicore__ inline HcclHandle AllGather(GM_ADDR sendBuf, GM_ADDR recvBuf, uint64_t sendCount, HcclDataType dataType,
        uint64_t strideCount, uint8_t repeat = 1);

    /*!
     * @class Hccl
     * @brief 集合通信算子ReduceScatter的任务下发接口（任务内容为：将所有rank的输入相加（或其他操作）后，再把结果按照rank编号均匀
     *        分散的到各个rank的输出buffer，每个进程拿到其他进程1/ranksize份的数据进行归约操作)，返回该任务的标识handleId给用户
     * @tparam commit: true/false 该参数设为true时，表示同步通知服务端可以执行对应的ReduceScatter通信任务
     * @param [out] recvBuf: 目的数据buffer地址，集合通信结果输出至此buffer中
     * @param [in] sendBuf: 源数据buffer地址
     * @param [in] recvCount: 参与ReduceScatter操作的recvBuf的数据个数，sendBuf的数据个数则等于recvCount * rank size。
     * @param [in] dataType: ReduceScatter操作的数据类型
     * @param [in] op: Reduce的操作类型，目前支持sum、prod、max、min等操作类型
     * @param [in] strideCount: 当多个rank的数据拼接到一个rank时，由于可能是多轮切分，因此需要strideCount参数告诉每一轮切分的数据个数
     * @param [in] repeat: 表明本通信任务需要执行的次数，默认为1，repeat必须>=1。
     *                     当repeat>1时，偏移数据在服务端由count*sizeof(dataType)算出，且只能用于连续场景
     * @return 返回该任务的标识handleId给用户，handleId >= 0。当本接口调用失败时，返回 -1。
     */
    template <bool commit = false>
    __aicore__ inline HcclHandle ReduceScatter(GM_ADDR sendBuf, GM_ADDR recvBuf, uint64_t recvCount,
        HcclDataType dataType, HcclReduceOp op, uint64_t strideCount, uint8_t repeat = 1);

public:  // 集合通信客户端控制函数，包括初始化/通知服务端/等待服务端等接口
    /*!
     * @class Hccl
     * @brief Hccl客户端初始化接口
     * @param [in] context: 通信context信息，包含rankDim，rankID等信息，对应数据结构HcclCombineOpParam
     */
    __aicore__ inline void Init(GM_ADDR context);

    /*!
     * @class Hccl
     * @brief 每调用一次本接口，则通知服务端可以正式执行handleId对应的任务一次，本接口的调用次数应该与Prepare的repeat次数一致
     * @param [in] handleId: 对应通信任务的标识ID，只能使用Prepare原语对应接口的返回值
     */
    __aicore__ inline void Commit(HcclHandle handleId);

    /*!
     * @class Hccl
     * @brief 阻塞Aicore，等待handleId对应的通信任务完成1次处理，本接口的调用次数应该与Prepare的repeat次数一致
     *        handleId调用Wait接口的顺序，须和Prepare接口一致
     * @param [in] handleId: 对应通信任务的标识ID，只能使用Prepare原语对应接口的返回值
     * @return 0表示成功，-1表示失败
     */
    __aicore__ inline int32_t Wait(HcclHandle handleId);

    /*!
     * @class Hccl
     * @brief 查询handleId对应的通信任务完成的轮次，最多返回repeat轮
     * @param [in] handleId: 对应通信任务的标识ID，只能使用Prepare原语对应接口的返回值
     * @return 返回handleId对应的通信任务的执行次数，最多repeat轮。当执行异常时，返回-1
     */
    __aicore__ inline int32_t Query(HcclHandle handleId);

    /*!
     * @class Hccl
     * @brief 通知服务端后续无通信任务，执行完通信任务即可退出
     */
    __aicore__ inline void Finalize();

    /*!
     * @class Hccl
     * @brief Hccl客户端的析构函数，析构时清空服务端的消息区(finishedTurnCnt)
     */
    __aicore__ inline ~Hccl();

public:
    /*!
     * @class Hccl
     * @brief 获取卡间通信数据WindowsIn起始地址，可用来直接作为作为计算的输入输出，减少拷贝
     * @param [in] rankId: 待查询的卡的Id
     * @return 返回对应卡的卡间通信数据WindowsIn起始地址
     */
    __aicore__ inline GM_ADDR GetWindowsInAddr(uint32_t rankId);

    /*!
     * @class Hccl
     * @brief 获取卡间通信数据WindowsOut起始地址，可用来直接作为作为计算的输入输出，减少拷贝
     * @param [in] rankId: 待查询的卡的Id
     * @return 返回对应卡的卡间通信数据WindowsOut起始地址
     */
    __aicore__ inline GM_ADDR GetWindowsOutAddr(uint32_t rankId);

    /*!
     * @class Hccl
     * @brief 获取本卡RankId
     * @return 返回本卡RankId
     */
    __aicore__ inline uint32_t GetRankId();

    /*!
     * @class Hccl
     * @brief 获取通信域内卡的数量
     * @return 获取通信域内卡的数量
     */
    __aicore__ inline uint32_t GetRankDim();

private:
    /*!
     * @class Hccl
     * @brief 当前版本只检查使用的消息数量不得超过消息总数(HCCL_MSG_CNT)
     * @return true:消息区可以正常使用  false:消息区已满，暂无法使用
     */
    __aicore__ inline bool CheckMsgListValid();

    /*!
     * @class Hccl
     * @brief 在aicore退出之前，清空消息区的finishedTurnCnt，保证下次launch的时候不会错误启动
     */
    __aicore__ inline void ResetFinishedTurnCnt();

    /*!
     * @class Hccl
     * @brief 各Prepare原语对应接口的通用实现
     * @param [in] necessaryHcclMsg: 各Prepare原语对应接口的必要入参，也即HcclMsg中需要用户自定义的参数、字段
     * @return 返回该任务的标识handleId给用户，handleId >= 0。当本接口调用失败时，返回 -1。
     */
    template<bool commit = false>
    __aicore__ inline HcclHandle CommonPrepareImpl(const NecessaryHcclMsg& necessaryHcclMsg);

private:
    __gm__ HcclCombineOpParam *hcclContext;
    // 消息区
    __gm__ HcclMsgArea *hcclMsgArea;

    // 暂时不开放用户设置，写死为0，表示只有0核会通过修改消息区与服务端交互，而其他核只能读取消息区，不影响服务端
    // Prepare系列接口、Commit接口、~Hccl析构函数中，只有0核能够写hcclMsgArea的字段
    // 910B上，除了Init，其他接口应该根据实际情况选择在cube或者vector上运行
    int32_t workBlockIdx;
    // 当前正在工作的消息在消息队列中的位置，从0开始自增。
    // Prepare/BatchPrepare/Finalize/InterHcclGroupSync(后续版本支持)都会且只占用一条消息
    uint8_t curMsgPosition;
    // 从0开始自增。一次prepare，不管repeat是多少，handle都只增加1
    // 一个HandleId对应的Commit和Wait调用次数应该与其Preapre的repeat参数值相等
    HcclHandle curHandleId;
    int8_t handleIdMsgPosition[HCCL_MAX_HANDLE_ID];
    // handleId任务对应的Commit次数
    uint8_t handleIdCommitTurnCnt[HCCL_MAX_HANDLE_ID];
    uint8_t handleIdRepeat[HCCL_MAX_HANDLE_ID];
    uint8_t handleIdWaitCallNum[HCCL_MAX_HANDLE_ID];
    bool isInited = false;
};
}  // namespace AscendC

#include "../../impl/hccl/hccl_v220_impl.h"
#endif  // __CCE_AICORE__ == 220
#endif  // LIB_HCCL_HCCL_H
